
#include "Tutorial_SSAOGameState.h"
#include "CameraController.h"
#include "GraphicsSystem.h"

#include "OgreBitwise.h"
#include "OgreItem.h"
#include "OgreSceneManager.h"

#include "OgreMesh2.h"
#include "OgreMeshManager.h"
#include "OgreMeshManager2.h"

#include "OgreCamera.h"
#include "OgreWindow.h"

#include "OgreHlmsPbsDatablock.h"
#include "OgreHlmsSamplerblock.h"

#include "OgreHlmsManager.h"
#include "OgreHlmsPbs.h"
#include "OgreRoot.h"
#include "OgreStagingTexture.h"
#include "OgreTextureBox.h"
#include "OgreTextureFilters.h"
#include "OgreTextureGpuManager.h"

#include "Vao/OgreVaoManager.h"
#include "Vao/OgreVertexArrayObject.h"

#include "OgreMaterialManager.h"
#include "OgrePass.h"
#include "OgreTechnique.h"

#include "Ogre.h"
#include "OgreResourceGroupManager.h"

using namespace Demo;

namespace Demo
{
    Tutorial_SSAOGameState::Tutorial_SSAOGameState( const Ogre::String &helpDescription ) :
        TutorialGameState( helpDescription ),
        mAnimateObjects( true ),
        mNumSpheres( 0 ),
        mKernelRadius( 1.0f ),
        mPowerScale( 1.5f )
    {
        memset( mSceneNode, 0, sizeof( mSceneNode ) );
    }
    //-----------------------------------------------------------------------------------
    void Tutorial_SSAOGameState::createScene01()
    {
        Ogre::SceneManager *sceneManager = mGraphicsSystem->getSceneManager();

        // Lets make some environment to look at
        // We use the same code as PBS demo to create meshes

        Ogre::v1::MeshPtr planeMeshV1 = Ogre::v1::MeshManager::getSingleton().createPlane(
            "Plane v1", Ogre::ResourceGroupManager::DEFAULT_RESOURCE_GROUP_NAME,
            Ogre::Plane( Ogre::Vector3::UNIT_Y, 1.0f ), 50.0f, 50.0f, 1, 1, true, 1, 4.0f, 4.0f,
            Ogre::Vector3::UNIT_Z, Ogre::v1::HardwareBuffer::HBU_STATIC,
            Ogre::v1::HardwareBuffer::HBU_STATIC );

        Ogre::MeshPtr planeMesh = Ogre::MeshManager::getSingleton().createByImportingV1(
            "Plane", Ogre::ResourceGroupManager::DEFAULT_RESOURCE_GROUP_NAME, planeMeshV1.get(), true,
            true, true );

        {
            Ogre::Item *item = sceneManager->createItem( planeMesh, Ogre::SCENE_DYNAMIC );
            item->setDatablock( "Marble" );
            Ogre::SceneNode *sceneNode = sceneManager->getRootSceneNode( Ogre::SCENE_DYNAMIC )
                                             ->createChildSceneNode( Ogre::SCENE_DYNAMIC );
            sceneNode->setPosition( 0, -1, 0 );
            sceneNode->attachObject( item );

            // Change the addressing mode of the roughness map to wrap via code.
            // Detail maps default to wrap, but the rest to clamp.
            assert( dynamic_cast<Ogre::HlmsPbsDatablock *>( item->getSubItem( 0 )->getDatablock() ) );
            Ogre::HlmsPbsDatablock *datablock =
                static_cast<Ogre::HlmsPbsDatablock *>( item->getSubItem( 0 )->getDatablock() );
            // Make a hard copy of the sampler block
            Ogre::HlmsSamplerblock samplerblock( *datablock->getSamplerblock( Ogre::PBSM_ROUGHNESS ) );
            samplerblock.mU = Ogre::TAM_WRAP;
            samplerblock.mV = Ogre::TAM_WRAP;
            samplerblock.mW = Ogre::TAM_WRAP;
            // Set the new samplerblock. The Hlms system will
            // automatically create the API block if necessary
            datablock->setSamplerblock( Ogre::PBSM_ROUGHNESS, samplerblock );
        }

        {
            mNumSpheres = 0;
            Ogre::HlmsManager *hlmsManager = mGraphicsSystem->getRoot()->getHlmsManager();

            assert( dynamic_cast<Ogre::HlmsPbs *>( hlmsManager->getHlms( Ogre::HLMS_PBS ) ) );

            Ogre::HlmsPbs *hlmsPbs =
                static_cast<Ogre::HlmsPbs *>( hlmsManager->getHlms( Ogre::HLMS_PBS ) );

            const int numX = 8;
            const int numZ = 8;

            const float armsLengthSphere = 2.0f;
            const float startX = ( numX - 1 ) / 2.0f;
            const float startZ = ( numZ - 1 ) / 2.0f;

            Ogre::Root *root = mGraphicsSystem->getRoot();
            Ogre::TextureGpuManager *textureMgr = root->getRenderSystem()->getTextureGpuManager();

            for( int x = 0; x < numX; ++x )
            {
                for( int z = 0; z < numZ; ++z )
                {
                    Ogre::String datablockName =
                        "Test" + Ogre::StringConverter::toString( mNumSpheres++ );
                    Ogre::HlmsPbsDatablock *datablock = static_cast<Ogre::HlmsPbsDatablock *>(
                        hlmsPbs->createDatablock( datablockName, datablockName, Ogre::HlmsMacroblock(),
                                                  Ogre::HlmsBlendblock(), Ogre::HlmsParamVec() ) );

                    Ogre::TextureGpu *texture = textureMgr->createOrRetrieveTexture(
                        "SaintPetersBasilica.dds", Ogre::GpuPageOutStrategy::Discard,
                        Ogre::TextureFlags::PrefersLoadingFromFileAsSRGB, Ogre::TextureTypes::TypeCube,
                        Ogre::ResourceGroupManager::AUTODETECT_RESOURCE_GROUP_NAME,
                        Ogre::TextureFilter::TypeGenerateDefaultMipmaps );

                    datablock->setTexture( Ogre::PBSM_REFLECTION, texture );
                    datablock->setDiffuse( Ogre::Vector3( 0.0f, 1.0f, 0.0f ) );

                    datablock->setRoughness(
                        std::max( 0.02f, float( x ) / std::max( 1.0f, (float)( numX - 1 ) ) ) );
                    datablock->setFresnel(
                        Ogre::Vector3( float( z ) / std::max( 1.0f, (float)( numZ - 1 ) ) ), false );

                    std::string meshName;
                    float meshScale = 1.0f;
                    if( x == z )
                    {
                        meshName = "Sphere1000.mesh";
                        meshScale = 2.0f;
                    }
                    else
                    {
                        meshName = "Sphere1000.mesh";
                        meshScale = 3.0f;
                    }

                    Ogre::Item *item = sceneManager->createItem(
                        meshName, Ogre::ResourceGroupManager::AUTODETECT_RESOURCE_GROUP_NAME,
                        Ogre::SCENE_DYNAMIC );
                    item->setDatablock( datablock );
                    item->setVisibilityFlags( 0x000000002 );

                    if( x != z )
                        item->setDatablock( "Marble" );

                    Ogre::SceneNode *sceneNode = sceneManager->getRootSceneNode( Ogre::SCENE_DYNAMIC )
                                                     ->createChildSceneNode( Ogre::SCENE_DYNAMIC );
                    sceneNode->setPosition(
                        Ogre::Vector3( armsLengthSphere * Ogre::Real( x ) - startX,
                                       0.0f + Ogre::Real( numZ - z ) * 0.5f,
                                       armsLengthSphere * Ogre::Real( z ) - startZ ) );
                    sceneNode->setScale( Ogre::Vector3( meshScale ) );
                    sceneNode->attachObject( item );
                }
            }
        }

        Ogre::SceneNode *rootNode = sceneManager->getRootSceneNode();

        Ogre::Light *light = sceneManager->createLight();
        Ogre::SceneNode *lightNode = rootNode->createChildSceneNode();
        lightNode->attachObject( light );
        light->setPowerScale( 1.0f );
        light->setType( Ogre::Light::LT_DIRECTIONAL );
        light->setDirection( Ogre::Vector3( -1, -1, -1 ).normalisedCopy() );

        mLightNodes[0] = lightNode;

        sceneManager->setAmbientLight( Ogre::ColourValue( 0.3f, 0.5f, 0.7f ) * 0.1f * 0.75f,
                                       Ogre::ColourValue( 0.6f, 0.45f, 0.3f ) * 0.065f * 0.75f,
                                       -light->getDirection() + Ogre::Vector3::UNIT_Y * 0.2f );

        light = sceneManager->createLight();
        lightNode = rootNode->createChildSceneNode();
        lightNode->attachObject( light );
        light->setDiffuseColour( 0.8f, 0.4f, 0.2f );  // Warm
        light->setSpecularColour( 0.8f, 0.4f, 0.2f );
        light->setPowerScale( Ogre::Math::PI );
        light->setType( Ogre::Light::LT_SPOTLIGHT );
        lightNode->setPosition( -10.0f, 10.0f, 10.0f );
        light->setDirection( Ogre::Vector3( 1, -1, -1 ).normalisedCopy() );
        light->setAttenuationBasedOnRadius( 10.0f, 0.01f );

        mLightNodes[1] = lightNode;

        light = sceneManager->createLight();
        lightNode = rootNode->createChildSceneNode();
        lightNode->attachObject( light );
        light->setDiffuseColour( 0.2f, 0.4f, 0.8f );  // Cold
        light->setSpecularColour( 0.2f, 0.4f, 0.8f );
        light->setPowerScale( Ogre::Math::PI );
        light->setType( Ogre::Light::LT_SPOTLIGHT );
        lightNode->setPosition( 10.0f, 10.0f, -10.0f );
        light->setDirection( Ogre::Vector3( -1, -1, 1 ).normalisedCopy() );
        light->setAttenuationBasedOnRadius( 10.0f, 0.01f );

        mLightNodes[2] = lightNode;

        mCameraController = new CameraController( mGraphicsSystem, false );

        //---------------------------------------------------------------------------------

        // We need to create SSAO kernel samples and noise texture
        // Generate kernel samples first
        float kernelSamples[64][4];
        for( size_t i = 0; i < 64u; ++i )
        {
            //            Ogre::Vector3 sample( 10, 10, 10 );
            //            while( sample.squaredLength() > 1.0f )
            //            {
            //                sample = Ogre::Vector3( Ogre::Math::RangeRandom(  -1.0f, 1.0f ),
            //                                        Ogre::Math::RangeRandom(  -1.0f, 1.0f ),
            //                                        Ogre::Math::RangeRandom(  0.01f, 1.0f ) );
            ////                sample = Ogre::Vector3( Ogre::Math::RangeRandom(  -0.1f, 0.1f ),
            ////                                        Ogre::Math::RangeRandom(  -0.1f, 0.1f ),
            ////                                        Ogre::Math::RangeRandom(  0.5f, 1.0f ) );
            //            }
            Ogre::Vector3 sample = Ogre::Vector3( Ogre::Math::RangeRandom( -1.0f, 1.0f ),
                                                  Ogre::Math::RangeRandom( -1.0f, 1.0f ),
                                                  Ogre::Math::RangeRandom( 0.0f, 1.0f ) );

            sample.normalise();

            float scale = (float)i / 64.0f;
            scale = Ogre::Math::lerp( 0.3f, 1.0f, scale * scale );
            sample = sample * scale;

            kernelSamples[i][0] = sample.x;
            kernelSamples[i][1] = sample.y;
            kernelSamples[i][2] = sample.z;
            kernelSamples[i][3] = 1.0f;
        }

        // Next generate noise texture
        Ogre::Root *root = mGraphicsSystem->getRoot();
        Ogre::TextureGpuManager *textureManager = root->getRenderSystem()->getTextureGpuManager();
        Ogre::TextureGpu *noiseTexture = textureManager->createTexture(
            "noiseTexture", Ogre::GpuPageOutStrategy::SaveToSystemRam, 0, Ogre::TextureTypes::Type2D );
        noiseTexture->setResolution( 2u, 2u );
        noiseTexture->setPixelFormat( Ogre::PFG_RGBA8_SNORM );
        noiseTexture->_transitionTo( Ogre::GpuResidency::Resident, (Ogre::uint8 *)0 );
        noiseTexture->_setNextResidencyStatus( Ogre::GpuResidency::Resident );

        Ogre::StagingTexture *stagingTexture =
            textureManager->getStagingTexture( 2u, 2u, 1u, 1u, Ogre::PFG_RGBA8_SNORM );
        stagingTexture->startMapRegion();
        Ogre::TextureBox texBox = stagingTexture->mapRegion( 2u, 2u, 1u, 1u, Ogre::PFG_RGBA8_SNORM );

        for( size_t j = 0; j < texBox.height; ++j )
        {
            for( size_t i = 0; i < texBox.width; ++i )
            {
                Ogre::Vector3 noise = Ogre::Vector3( Ogre::Math::RangeRandom( -1.0f, 1.0f ),
                                                     Ogre::Math::RangeRandom( -1.0f, 1.0f ), 0.0f );
                noise.normalise();

                Ogre::int8 *pixelData = reinterpret_cast<Ogre::int8 *>( texBox.at( i, j, 0 ) );
                pixelData[0] = Ogre::Bitwise::floatToSnorm8( noise.x );
                pixelData[1] = Ogre::Bitwise::floatToSnorm8( noise.y );
                pixelData[2] = Ogre::Bitwise::floatToSnorm8( noise.z );
                pixelData[3] = Ogre::Bitwise::floatToSnorm8( 1.0f );
            }
        }

        stagingTexture->stopMapRegion();
        stagingTexture->upload( texBox, noiseTexture, 0, 0, 0 );
        textureManager->removeStagingTexture( stagingTexture );
        stagingTexture = 0;

        //---------------------------------------------------------------------------------
        // Get GpuProgramParametersSharedPtr to set uniforms that we need
        Ogre::MaterialPtr material =
            std::static_pointer_cast<Ogre::Material>( Ogre::MaterialManager::getSingleton().load(
                "SSAO/HS", Ogre::ResourceGroupManager::AUTODETECT_RESOURCE_GROUP_NAME ) );

        Ogre::Pass *pass = material->getTechnique( 0 )->getPass( 0 );
        mSSAOPass = pass;
        Ogre::GpuProgramParametersSharedPtr psParams = pass->getFragmentProgramParameters();

        // Lets set uniforms for shader
        // Set texture uniform for noise
        Ogre::TextureUnitState *noiseTextureState = pass->getTextureUnitState( "noiseTexture" );
        noiseTextureState->setTexture( noiseTexture );

        // Reconstruct position from depth. Position is needed in SSAO
        // We need to set the parameters based on camera to the
        // shader so that the un-projection works as expected
        Ogre::Camera *camera = mGraphicsSystem->getCamera();
        Ogre::Vector2 projectionAB = camera->getProjectionParamsAB();
        // The division will keep "linearDepth" in the shader in the [0; 1] range.
        projectionAB.y /= camera->getFarClipDistance();
        psParams->setNamedConstant( "projectionParams", projectionAB );

        // Set other uniforms
        psParams->setNamedConstant( "kernelRadius", mKernelRadius );
        psParams->setNamedConstant(
            "noiseScale",
            Ogre::Vector2(
                ( Ogre::Real( mGraphicsSystem->getRenderWindow()->getWidth() ) * 0.5f ) / 2.0f,
                ( Ogre::Real( mGraphicsSystem->getRenderWindow()->getHeight() ) * 0.5f ) / 2.0f ) );
        psParams->setNamedConstant( "invKernelSize", 1.0f / 64.0f );
        psParams->setNamedConstant( "sampleDirs", (float *)kernelSamples, 64, 4 );

        // Set blur shader uniforms
        Ogre::MaterialPtr materialBlurH =
            std::static_pointer_cast<Ogre::Material>( Ogre::MaterialManager::getSingleton().load(
                "SSAO/BlurH", Ogre::ResourceGroupManager::AUTODETECT_RESOURCE_GROUP_NAME ) );

        Ogre::Pass *passBlurH = materialBlurH->getTechnique( 0 )->getPass( 0 );
        Ogre::GpuProgramParametersSharedPtr psParamsBlurH = passBlurH->getFragmentProgramParameters();
        psParamsBlurH->setNamedConstant( "projectionParams", projectionAB );

        Ogre::MaterialPtr materialBlurV =
            std::static_pointer_cast<Ogre::Material>( Ogre::MaterialManager::getSingleton().load(
                "SSAO/BlurV", Ogre::ResourceGroupManager::AUTODETECT_RESOURCE_GROUP_NAME ) );

        Ogre::Pass *passBlurV = materialBlurV->getTechnique( 0 )->getPass( 0 );
        Ogre::GpuProgramParametersSharedPtr psParamsBlurV = passBlurV->getFragmentProgramParameters();
        psParamsBlurV->setNamedConstant( "projectionParams", projectionAB );

        // Set apply shader uniforms
        Ogre::MaterialPtr materialApply =
            std::static_pointer_cast<Ogre::Material>( Ogre::MaterialManager::getSingleton().load(
                "SSAO/Apply", Ogre::ResourceGroupManager::AUTODETECT_RESOURCE_GROUP_NAME ) );

        Ogre::Pass *passApply = materialApply->getTechnique( 0 )->getPass( 0 );
        mApplyPass = passApply;
        Ogre::GpuProgramParametersSharedPtr psParamsApply = passApply->getFragmentProgramParameters();
        psParamsApply->setNamedConstant( "powerScale", mPowerScale );

        TutorialGameState::createScene01();
    }
    //-----------------------------------------------------------------------------------
    void Tutorial_SSAOGameState::update( float timeSinceLast )
    {
        Ogre::GpuProgramParametersSharedPtr psParams = mSSAOPass->getFragmentProgramParameters();
        Ogre::Camera *camera = mGraphicsSystem->getCamera();
#if OGRE_NO_VIEWPORT_ORIENTATIONMODE == 0
        // We don't render to render window directly, thus we need to get the projection
        // matrix with phone orientation disable when calculating SSAO
        camera->setOrientationMode( Ogre::OR_DEGREE_0 );
#endif
        psParams->setNamedConstant( "projection", camera->getProjectionMatrix() );
        psParams->setNamedConstant( "kernelRadius", mKernelRadius );

        Ogre::GpuProgramParametersSharedPtr psParamsApply = mApplyPass->getFragmentProgramParameters();
        psParamsApply->setNamedConstant( "powerScale", mPowerScale );

        TutorialGameState::update( timeSinceLast );
    }
    //-----------------------------------------------------------------------------------
    void Tutorial_SSAOGameState::generateDebugText( float timeSinceLast, Ogre::String &outText )
    {
        TutorialGameState::generateDebugText( timeSinceLast, outText );

        if( mDisplayHelpMode == 1 )
        {
            outText += "\nPress F5/F6 to increase/reduce kernel radius. ";
            outText += "[" + Ogre::StringConverter::toString( mKernelRadius ) + "]";
            outText += "\nPress F7/F8 to increase/reduce power scale. ";
            outText += "[" + Ogre::StringConverter::toString( mPowerScale ) + "]";
        }
    }
    //-----------------------------------------------------------------------------------
    void Tutorial_SSAOGameState::keyReleased( const SDL_KeyboardEvent &arg )
    {
        if( ( arg.keysym.mod & ~( KMOD_NUM | KMOD_CAPS ) ) != 0 )
        {
            TutorialGameState::keyReleased( arg );
            return;
        }

        if( arg.keysym.sym == SDLK_F2 )
        {
            mAnimateObjects = !mAnimateObjects;
        }
        else if( arg.keysym.sym == SDLK_F5 )
        {
            mKernelRadius = mKernelRadius - 0.1f;
            if( mKernelRadius < 0.05f )
                mKernelRadius = 0.05f;
        }
        else if( arg.keysym.sym == SDLK_F6 )
        {
            mKernelRadius = mKernelRadius + 0.1f;
            if( mKernelRadius > 10.0f )
                mKernelRadius = 10.0f;
        }
        else if( arg.keysym.sym == SDLK_F7 )
        {
            mPowerScale = mPowerScale - 0.25f;
        }
        else if( arg.keysym.sym == SDLK_F8 )
        {
            mPowerScale = mPowerScale + 0.25f;
        }
        else
        {
            TutorialGameState::keyReleased( arg );
        }
    }
}  // namespace Demo
